import json
import os

import cv2
import joblib
import numpy as np
from tkinter import Tk, Button, filedialog
import matplotlib.pyplot as plt
from skimage.feature import hog

# 加载模型
model_path = "model/svm_model_1221.joblib"  # 模型路径
label_map = {0: 'Beech Mushroom', 1: 'Black Trumpet Mushroom', 2: 'Button Mushroom', 3: 'Chanterelle Mushroom',
             4: 'Crimini Mushroom', 5: 'Enoki Mushroom', 6: 'Hedgehog Mushroom', 7: 'King Oyster Mushroom',
             8: 'Maitake Mushroom', 9: 'Morel Mushroom', 10: 'Oyster Mushroom', 11: 'Porcini Mushroom',
             12: 'Portobello Mushroom', 13: 'Shiitake Mushroom'}  # 标签映射
svm_model = joblib.load(model_path)


# 图像预处理函数：调整图像大小，灰度化，归一化
def preprocess_image(image, target_size=(128, 128)):
    image_resized = cv2.resize(image, target_size)
    gray_image = cv2.cvtColor(image_resized, cv2.COLOR_BGR2GRAY)
    normalized_image = gray_image / 255.0
    return normalized_image


# 提取图像的 HOG 特征
def extract_hog_features(images):
    hog_features = []
    for image in images:
        hog_feature = hog(image, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2), transform_sqrt=True)
        hog_features.append(hog_feature)
    return np.array(hog_features)


# 植物图像识别
def predict_plant(image, svm_model, label_map):
    preprocessed_image = preprocess_image(image)
    hog_feature = extract_hog_features([preprocessed_image])
    confidence = svm_model.decision_function(hog_feature)

    # 将置信度映射到 [0, 1] 的范围
    min_confidence = np.min(confidence)
    max_confidence = np.max(confidence)
    scaled_confidence = (confidence - min_confidence) / (max_confidence - min_confidence)
    print('scaled_confidence', scaled_confidence)
    prediction = svm_model.predict(hog_feature)
    plant_label = label_map[prediction[0]]
    return plant_label


def expected_img_name(file_path):
    # 获取文件的目录
    directory = os.path.dirname(file_path)
    # 获取最底层的文件夹名称
    last_folder_name = os.path.basename(directory)
    return last_folder_name

# 加载mushroom.json
def load_mushroom_data():
    try:
        with open('mushroom.json', 'r', encoding='utf-8') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading mushroom data: {e}")
        return []


# mushroom.json的列表中找到type_name对应一项
def find_mushroom_by_type(type_name, mushroom_data):
    return next(
        (item for item in mushroom_data if item["type_name"].lower() == type_name.lower()),
        None
    )


def load_image():
    root = Tk()
    root.withdraw()
    file_path = filedialog.askopenfilename(title="Select an Image")
    if file_path:
        image = cv2.imread(file_path)
        plant_label = predict_plant(image, svm_model, label_map)
        plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

        # 加载蘑菇数据
        # mushroom_data = load_mushroom_data()
        # type_chinese_name = find_mushroom_by_type(plant_label, mushroom_data)['type_chinese_name']
        # plt.title(f'Predicted Plant: {plant_label},chinese name {type_chinese_name}')

        expect_type_label = expected_img_name(file_path)
        # expect_type_chinese_name = find_mushroom_by_type(expect_type_label, mushroom_data)['type_chinese_name']

        plt.title(f'Predicted Plant: {plant_label},\n expect Plant: {expect_type_label}')
        plt.axis('off')
        plt.show()


# 创建GUI按钮
root = Tk()
root.title("Plant Image Recognition")
button = Button(root, text="Load Image", command=load_image)
button.pack(pady=20)
root.mainloop()
